"""
基于 CrewAI 的智能体协作流程调度器
"""
from crewai import Crew, Process
from agents.data_collector_agent import DataCollectorAgent
from agents.data_analyst_agent import DataAnalystAgent
from agents.decision_maker_agent import DecisionMakerAgent
from colorama import Fore, Style, init
from typing import Dict, Any, List
import time
import json
import os


class CrewAIStockFlow:
    """基于 CrewAI 的股票分析智能体协作流程"""
    
    def __init__(self):
        init(autoreset=True)
        print(Fore.CYAN + "🚀 初始化 CrewAI 智能体协作系统...")
        
        # 初始化智能体
        self.data_collector = DataCollectorAgent()
        self.data_analyst = DataAnalystAgent()
        self.decision_maker = DecisionMakerAgent()
        
        # 加载配置
        self.crew_config = self._load_crew_config()
        
        print(Fore.GREEN + "✅ CrewAI 智能体系统初始化完成")
    
    def _load_crew_config(self) -> Dict[str, Any]:
        """加载 Crew 配置"""
        try:
            config_path = os.path.join(os.path.dirname(__file__), "config.json")
            with open(config_path, "r", encoding="utf-8") as f:
                config = json.load(f)
                return config.get("crew_config", {})
        except Exception as e:
            print(Fore.YELLOW + f"⚠️ 加载配置失败，使用默认配置: {e}")
            return {
                "verbose": True,
                "memory": True,
                "max_iter": 5,
                "planning": False
            }
    
    def run(self, stock_code: str):
        """执行完整的 CrewAI 智能体协作流程"""
        print(Fore.CYAN + Style.BRIGHT + f"\n🎬 === CrewAI 智能体协作开始 ===")
        print(Fore.CYAN + f"📈 目标股票: {stock_code}")
        print(Fore.CYAN + f"🤖 智能体数量: 3 个")
        print(Fore.CYAN + f"⚡ 处理引擎: CrewAI")
        
        start_time = time.time()
        
        try:
            # 创建任务
            print(Fore.MAGENTA + f"\n🔄 [CrewAI] 创建智能体任务...")
            
            # 数据收集任务
            collector_tasks = self.data_collector.create_tasks(stock_code)
            
            # 获取最后一个收集任务作为分析的输入
            last_collection_task = collector_tasks[-1] if collector_tasks else None
            
            # 数据分析任务
            analyst_tasks = self.data_analyst.create_tasks(
                stock_code, 
                data_collection_task=last_collection_task
            )
            
            # 获取最后一个分析任务作为决策的输入
            last_analysis_task = analyst_tasks[-1] if analyst_tasks else None
            
            # 投资决策任务
            decision_tasks = self.decision_maker.create_tasks(
                stock_code,
                analysis_task=last_analysis_task
            )
            
            # 合并所有任务
            all_tasks = collector_tasks + analyst_tasks + decision_tasks
            
            print(Fore.GREEN + f"✅ 创建了 {len(all_tasks)} 个任务")
            
            # 创建 Crew
            print(Fore.MAGENTA + f"\n🔄 [CrewAI] 组建智能体团队 (Crew)...")
            
            crew = Crew(
                agents=[
                    self.data_collector.get_agent(),
                    self.data_analyst.get_agent(),
                    self.decision_maker.get_agent()
                ],
                tasks=all_tasks,
                verbose=self.crew_config.get("verbose", True),
                process=Process.sequential,  # 顺序执行
                memory=self.crew_config.get("memory", True),
                max_iter=self.crew_config.get("max_iter", 5)
            )
            
            print(Fore.GREEN + f"✅ Crew 组建完成，开始执行任务...")
            
            # 执行 Crew
            print(Fore.MAGENTA + f"\n🔄 [CrewAI] 启动智能体协作工作流...")
            result = crew.kickoff()
            
            end_time = time.time()
            processing_time = round(end_time - start_time, 2)
            
            # 显示结果
            self._display_results(result, processing_time)
            
            return result
            
        except Exception as e:
            print(Fore.RED + f"\n❌ CrewAI 工作流执行失败: {str(e)}")
            import traceback
            traceback.print_exc()
            return None
    
    def _display_results(self, result: Any, processing_time: float):
        """显示处理结果"""
        print(Fore.GREEN + f"\n✅ [CrewAI] 智能体协作工作流完成 (耗时: {processing_time}s)")
        
        print(Fore.CYAN + "\n📋 === 执行结果摘要 ===")
        
        # 显示结果
        if result:
            if isinstance(result, str):
                # 如果结果是字符串，直接显示
                print(Fore.WHITE + Style.BRIGHT + "\n📊 === 最终分析结果 ===")
                print(Fore.WHITE + result[:2000] + "..." if len(result) > 2000 else result)
            elif isinstance(result, dict):
                # 如果结果是字典，格式化显示
                print(Fore.WHITE + Style.BRIGHT + "\n📊 === 分析结果详情 ===")
                for key, value in result.items():
                    print(Fore.YELLOW + f"{key}:")
                    print(Fore.WHITE + f"  {str(value)[:500]}")
            else:
                # 其他类型，转换为字符串显示
                print(Fore.WHITE + str(result)[:2000])
        
        print(Fore.CYAN + Style.BRIGHT + f"\n🏁 === CrewAI 智能体协作结束 ===")
    
    def get_agent_info(self):
        """获取智能体信息"""
        agents_info = [
            self.data_collector.get_agent_info(),
            self.data_analyst.get_agent_info(),
            self.decision_maker.get_agent_info()
        ]
        
        print(Fore.CYAN + "\n🤖 === CrewAI 智能体信息 ===")
        for i, info in enumerate(agents_info, 1):
            print(f"{Fore.YELLOW}{i}. {info['name']}")
            print(f"   - 角色: {info['role']}")
            print(f"   - 目标: {info['goal']}")
            print(f"   - 模型: {info['model']}")
            print(f"   - 服务: {info['base_url']}")
            print(f"   - 工具: {', '.join(info['tools']) if info['tools'] else '无'}")
        
        return agents_info
    
    def run_simple_analysis(self, stock_code: str):
        """运行简化版分析（仅数据收集和基础分析）"""
        print(Fore.CYAN + f"\n🚀 运行简化分析模式...")
        
        try:
            # 仅创建数据收集和一个分析任务
            collector_tasks = self.data_collector.create_tasks(stock_code)[:1]
            
            # 创建简化的 Crew
            simple_crew = Crew(
                agents=[self.data_collector.get_agent()],
                tasks=collector_tasks,
                verbose=True,
                process=Process.sequential
            )
            
            result = simple_crew.kickoff()
            print(Fore.GREEN + "✅ 简化分析完成")
            return result
            
        except Exception as e:
            print(Fore.RED + f"❌ 简化分析失败: {e}")
            return None
